#!/usr/bin/env python
# -*- mode: python; encoding: utf-8; -*-
#
# Copyright (C) 2010 Jonathan Ellis
#
#'''"""HMMER v3 utils.
#
#- Output file may contain multiple records (separated by //).
#
#- Each record contains scores for complete sequences.
#
#- Each record contains info for multiple domain annotation (either for
#  sequence or model depending on whether hmmsearch or hmmscan used).

#- Each domain annotation potentially contains multiple domain alignments.

#should a hit combine the complete sequence info with the domain annotation?
#how? 
#
#Hit (complete, domain)
#NOTE FROM SAM
#minor change of the code
#1. fix bug when target sequence has name with special symbols
#"""'''
import os
import re
import sys
#from StringIO import StringIO
#from subprocess import Popen, PIPE

class CompleteSequence(object):
    def __init__(self, data):
        try:
            items = data.split()
            self.full_evalue = float(items[0])
            self.full_score  = float(items[1])
            self.full_bias   = float(items[2])
            self.best_evalue = float(items[3])
            self.best_score  = float(items[4])
            self.best_bias   = float(items[5])
            self.exp         = float(items[6])
            self.n           = int(items[7])
            self.sequence    = items[8] # or model for hmmscan
            self.name = items[8]
            # The description is not always provided.
            try:
                self.description = items[9]
            except IndexError:
                self.description = ''
        except:
            return

class Record(object):
    def __init__(self, record):
        self.query       = None
        self.accession   = None
        self.description = None
        
        header_flag      = True
        complete_flag    = False
        domain_flag      = False
        
        header_data      = []
        complete_data    = []
        domain_data      = []

        for line in record:
            # Scores for complete sequence(s) -- the final 's' seems to
            # depend on whether we're parsing hmmsearch (with) or hmmscan
            # (without) output.
            if line.startswith('Scores for complete sequence'):
                complete_flag = True
                header_flag = False
            if line.startswith('Domain annotation'):
                complete_flag = False
                domain_flag = True
            if line.startswith('Internal pipeline'):
                domain_flag = False
            if header_flag:
                header_data.append(line)
            if complete_flag:
                complete_data.append(line)
            if domain_flag:
                domain_data.append(line)

        self.hits = []
        
        self._parse_header(header_data)

        # It's possible there were no hits/domains; therefore, we only try
        # to parse the data if it's there.
        if self._have_domains(domain_data):
            self._parse_domain_annotation(domain_data)
        if self._have_hits(complete_data):
            self._parse_complete_sequences(complete_data)


    def _have_hits(self, complete_data):
        match = re.search(r'No hits detected that satisfy reporting thresholds',
                          ''.join(complete_data))
        return match is None

    def _have_domains(self, domain_data):
        match = re.search(r'No targets detected that satisfy reporting thresholds',
                          ''.join(domain_data))
        return match is None

    def _parse_header(self, header):
        for line in header:
            if line.startswith('Query:'):
                self.query = line.split()[1]

            if line.startswith('Accession:'):
                self.accession = line.lstrip('Accession:').strip()
            
            if line.startswith('Description:'):
                self.description = line.lstrip('Description:').strip()

    def _parse_complete_sequences(self, complete_data):
        # skip first four lines.
        complete_sequences = {}
        for line in complete_data[4:]:
            if not line.strip():
                continue
            if '------ inclusion threshold ------' in line:
                continue
            cs = CompleteSequence(line)
            complete_sequences[cs.name] = cs
        for hit in self.hits:
            cs = complete_sequences[hit.name]
            hit.full_evalue = cs.full_evalue
            hit.full_score  = cs.full_score 
            hit.full_bias   = cs.full_bias  
            hit.best_evalue = cs.best_evalue
            hit.best_score  = cs.best_score 
            hit.best_bias   = cs.best_bias  
            hit.exp         = cs.exp        
            hit.n           = cs.n
            hit.name        = cs.name
            hit.description = cs.description

    def _parse_domain_annotation(self, domain_data):
        first_domain = True
        a, b = [], []
        for line in domain_data[1:]:
            if line.startswith('>>'):
                if first_domain:
                    first_domain = False
                else:
                    a.append(b)
                    b = []
            b.append(line)
        a.append(b)
        self.hits = [Hit(x) for x in a]

    def __iter__(self):
        for hit in self.hits:
            yield hit

    def __len__(self):
        return len(self.hits)

    def __getitem__(self, x):
        return self.hits[x]


class Domain(object):
    def __init__(self, data, sequence_name):
        data = data.splitlines()
        match = re.search(r'domain\s+(\d+)', data[0])
        if match:
            self.domain = match.group(1)
        match = re.search(r'score: (.+?) bits;', data[0])
        if match:
            self.score = match.group(1)
        match = re.search(r'conditional E-value: (.+?)$', data[0])
        if match:
            self.evalue = float(match.group(1))

        posterior_prob  = []
        target_sequence = []
        consensus       = []
        hmm_sequence    = []

        for line in data[1:]:
            if line.endswith(' CS'):
                pass
            elif line.endswith(' RF'):
                pass
            elif line.endswith(' PP'):
                posterior_prob.append(line.rstrip('\r\n').rstrip('PP').strip())
            else:
                match = re.search(r'\s+(.+?)\s+(\d+)\s+(.+?)\s+(\d+)', line)
                #if the name of target sequence have special simbol (e.g. :) this will fail
                #was (re.search(r'\s+(\w+)\s+(\d+)\s+(.+?)\s+(\d+)', line))
                if match:
                    start_position = int(line.index(match.group(3)))
                    end_position = start_position + len(match.group(3))
                    if match.group(1) == sequence_name:
                        #print "match.group3", match.group(3)
                        target_sequence.append(match.group(3)) #was target_sequence
                    else:
                        hmm_sequence.append(match.group(3)) #was hmm_sequence
                else:
                    match = re.search(r'\s+(.+?)$', line)
                    if match:
                        # Consensus may start/end with whitespace so we need
                        # to be more precise with the extraction.
                        consensus.append(line[start_position:end_position])

        self.posterior_prob = ''.join(posterior_prob)
        self.target_sequence = ''.join(target_sequence)
        self.consensus = ''.join(consensus)
        self.hmm_sequence = ''.join(hmm_sequence)

    # name() is provided as a convience.
    @property
    def name(self):
        return self.parent.name

class Hit(object):
    # has info for complete sequence plus one or more domains
    def __init__(self, hit_data):
        try:
            self.name = hit_data[0].split()[1]
        except:
            return
        ## self.full_evalue = None
        ## self.full_score  = None
        ## self.full_bias   = None
        ## self.best_evalue = None
        ## self.best_score  = None
        ## self.best_bias   = None
        ## self.exp         = None
        ## self.n           = None
        self.domains = []
            
        flag = False
        d = {}
        for line in hit_data:
            match = re.search(r'\d+\s+[!?]', line)
            if match:
                (num, ex, score, bias, c_evalue, i_evalue,
                 hmm_from, hmm_to, hmm_bits,
                 ali_from, ali_to, ali_bits,
                 env_from, env_to, env_bits, acc) = line.split()

                d[num] = (score, bias, c_evalue, i_evalue,
                          hmm_from, hmm_to, hmm_bits,
                          ali_from, ali_to, ali_bits,
                          env_from, env_to, env_bits, acc)
        self.domains = self._parse_domains(hit_data)
        if self.domains:
            for domain in self.domains:
                domain.parent   = self
                n               = domain.domain
                domain.id       = domain.domain
                domain.score    = float(d[n][0])
                domain.bias     = float(d[n][1])
                domain.c_evalue = float(d[n][2])
                domain.i_evalue = float(d[n][3])
                domain.hmm_from = int(d[n][4])
                domain.hmm_to   = int(d[n][5])
                domain.hmm_bits = d[n][6]
                domain.ali_from = int(d[n][7])
                domain.ali_to   = int(d[n][8])
                domain.ali_bits = d[n][9]
                domain.env_from = int(d[n][10])
                domain.env_to   = int(d[n][11])
                domain.env_bits = d[n][12]
                domain.acc      = float(d[n][13])

    def _parse_domains(self, data):
        try:
            sequence_name = data[0].split()[1]
            domains = []
            first = True
            flag = False
            d = []
            for line in data:
                # Skip domain annotation header.
                if 'Alignments for each domain' in line:
                    flag = True
                    continue
                if flag:
                    if re.search(r'\s+==', line):
                        if first:
                            first = False
                        else:
                            domains.append(Domain('\n'.join(d), sequence_name))
                            d = []
                    d.append(line)
            domains.append(Domain('\n'.join(d), sequence_name))
        except:
            # If the --noali option was passed to hmm(search|scan) there
            # will be no alignments to pass, so return empty list.
            domains = []
        return domains

    def __iter__(self):
        for domain in self.domains:
            yield domain

    def __len__(self):
        return len(self.domains)

    def __getitem__(self, x):
        return self.domains[x]


def parse(handle):
    first_line = handle.readline()
    d = []
    if (first_line.startswith('# hmmsearch') or
        first_line.startswith('# hmmscan')):
        records = []
        record = []
        for line in handle:
            if line.startswith('#'):
                continue
            line = line.rstrip('\r\n')
            record.append(line)
            if line == '//':
                records.append(record)
                record = []
        for record in records:
            yield Record(record)
    else:
        raise ValueError('does not appear to be a HMMER file')


if __name__ == '__main__':
    for record in parse(sys.stdin):
        for hit in record:
            #print 'Hit', hit.name, hit.full_evalue
            for domain in hit:
                print "HI ITS ME"
                #print dir(domain)
                '''
                print 'ali_from', domain.ali_from
                print 'ali_to', domain.ali_to
                print 'ali bit', domain.ali_bits
                print 'bias', domain.bias
                print 'c_evalue', domain.c_evalue
                print 'i_evalue', domain.i_evalue
                '''
                #print 'consensus', domain.consensus
                #print 'domain', domain.domain
                #print 'score', domain.score
                #print 'posterior', domain.posterior_prob
                print domain.id
                print 'hmm seque', domain.hmm_sequence
                print 'target_se', domain.target_sequence
                '''
                print 'hmm from', domain.hmm_from
                print 'hmm_to', domain.hmm_to
                print 'hmm bits', domain.hmm_bits
                print 'env_from', domain.env_from
                print 'env_to', domain.env_to
                '''
